import tkinter as tk
from tkinter import filedialog, messagebox
from tkinter import ttk
import os

def read_srt(srt_file):
    """读取SRT文件，返回时间码和字幕文本的列表"""
    with open(srt_file, 'r', encoding='utf-8') as f:
        content = f.read().strip()
    
    blocks = content.split('\n\n')
    subtitles = []
    
    for block in blocks:
        lines = block.split('\n')
        if len(lines) >= 3:
            number = lines[0]
            timestamp = lines[1]
            text = '\n'.join(lines[2:])
            subtitles.append((number, timestamp, text))
    
    return subtitles

def read_translation(translation_file):
    """读取翻译文本文件，返回翻译文本列表"""
    with open(translation_file, 'r', encoding='utf-8') as f:
        translations = [line.strip() for line in f if line.strip()]
    return translations

def create_new_srt(subtitles, translations, output_srt, output_txt=None):
    """创建新的SRT文件和TXT文件"""
    if len(subtitles) != len(translations):
        raise ValueError("字幕数量与翻译数量不匹配！")
    
    # 写入SRT文件
    with open(output_srt, 'w', encoding='utf-8') as f:
        for (number, timestamp, _), translation in zip(subtitles, translations):
            f.write(f"{number}\n")
            f.write(f"{timestamp}\n")
            f.write(f"{translation}\n\n")
    
    # 如果指定了TXT输出路径，则同时写入TXT文件
    if output_txt:
        with open(output_txt, 'w', encoding='utf-8') as f:
            for (_, _, original), translation in zip(subtitles, translations):
                f.write(f"原文: {original}\n")
                f.write(f"译文: {translation}\n")
                f.write("-" * 50 + "\n")

def count_srt_lines(srt_file):
    """统计SRT文件中的字幕条数"""
    try:
        with open(srt_file, 'r', encoding='utf-8') as f:
            content = f.read().strip()
        blocks = content.split('\n\n')
        return len(blocks)
    except Exception:
        return 0

def count_txt_lines(txt_file):
    """统计TXT文件中的非空行数"""
    try:
        with open(txt_file, 'r', encoding='utf-8') as f:
            lines = [line.strip() for line in f if line.strip()]
        return len(lines)
    except Exception:
        return 0

class SrtTranslatorGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("SRT字幕翻译替换工具")
        self.root.geometry("800x600")  # 增加窗口大小以适应预览区域
        
        # 设置主题样式
        style = ttk.Style()
        style.theme_use('clam')  # 使用 'clam' 主题，也可以尝试 'alt', 'default', 'classic'
        
        # 创建主框架并设置权重
        main_frame = ttk.Frame(root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        root.columnconfigure(0, weight=1)
        root.rowconfigure(0, weight=1)
        
        # 配置主框架的列权重
        main_frame.columnconfigure(1, weight=1)
        
        # SRT文件选择
        ttk.Label(main_frame, text="SRT文件:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.srt_path = tk.StringVar()
        ttk.Entry(main_frame, textvariable=self.srt_path, width=50).grid(row=0, column=1, padx=5, sticky=(tk.W, tk.E))
        ttk.Button(main_frame, text="浏览", command=self.select_srt, width=10).grid(row=0, column=2, padx=5)
        
        # 翻译文件选择
        ttk.Label(main_frame, text="翻译文件:").grid(row=1, column=0, sticky=tk.W, pady=5)
        self.trans_path = tk.StringVar()
        ttk.Entry(main_frame, textvariable=self.trans_path, width=50).grid(row=1, column=1, padx=5, sticky=(tk.W, tk.E))
        ttk.Button(main_frame, text="浏览", command=self.select_translation, width=10).grid(row=1, column=2, padx=5)
        
        # 输出文件选择
        ttk.Label(main_frame, text="输出文件:").grid(row=2, column=0, sticky=tk.W, pady=5)
        self.output_path = tk.StringVar()
        ttk.Entry(main_frame, textvariable=self.output_path, width=50).grid(row=2, column=1, padx=5, sticky=(tk.W, tk.E))
        ttk.Button(main_frame, text="浏览", command=self.select_output, width=10).grid(row=2, column=2, padx=5)
        
        # 添加导出TXT选项
        self.export_txt = tk.BooleanVar(value=True)
        ttk.Checkbutton(
            main_frame, 
            text="同时导出TXT对照文件", 
            variable=self.export_txt
        ).grid(row=3, column=0, columnspan=3, pady=5)
        
        # 在状态标签之前添加行数统计标签
        self.count_frame = ttk.Frame(main_frame)
        self.count_frame.grid(row=4, column=0, columnspan=3, pady=5)
        
        self.srt_count_var = tk.StringVar(value="SRT字幕数: 0")
        self.txt_count_var = tk.StringVar(value="TXT行数: 0")
        
        ttk.Label(self.count_frame, textvariable=self.srt_count_var).pack(side=tk.LEFT, padx=10)
        ttk.Label(self.count_frame, textvariable=self.txt_count_var).pack(side=tk.LEFT, padx=10)
        
        # 转换按钮 - 调整行号
        button_frame = ttk.Frame(main_frame)
        button_frame.grid(row=5, column=0, columnspan=3, pady=10)
        convert_button = ttk.Button(button_frame, text="开始转换", command=self.convert, width=20)
        convert_button.pack()
        
        # 状态标签 - 调整行号
        self.status_var = tk.StringVar()
        self.status_var.set("准备就绪")
        status_label = ttk.Label(main_frame, textvariable=self.status_var)
        status_label.grid(row=6, column=0, columnspan=3, pady=5)
        
        # 添加预览区域
        preview_frame = ttk.Frame(main_frame)
        preview_frame.grid(row=7, column=0, columnspan=3, sticky=(tk.W, tk.E, tk.N, tk.S), pady=5)
        
        # 配置预览区域的权重
        preview_frame.columnconfigure(0, weight=1)
        preview_frame.columnconfigure(1, weight=1)
        preview_frame.rowconfigure(0, weight=0)
        preview_frame.rowconfigure(1, weight=1)
        
        # SRT预览标题
        ttk.Label(preview_frame, text="SRT文件预览:").grid(row=0, column=0, sticky=tk.W, pady=2)
        ttk.Label(preview_frame, text="翻译文件预览:").grid(row=0, column=1, sticky=tk.W, pady=2)
        
        # SRT预览文本框
        self.srt_preview = tk.Text(preview_frame, height=10, width=40)
        self.srt_preview.grid(row=1, column=0, padx=5, sticky=(tk.W, tk.E, tk.N, tk.S))
        srt_scrollbar = ttk.Scrollbar(preview_frame, orient=tk.VERTICAL, command=self.srt_preview.yview)
        srt_scrollbar.grid(row=1, column=0, sticky=(tk.N, tk.S, tk.E))
        self.srt_preview.configure(yscrollcommand=srt_scrollbar.set)
        self.srt_preview.configure(state='disabled')
        
        # 翻译文件预览文本框
        self.trans_preview = tk.Text(preview_frame, height=10, width=40)
        self.trans_preview.grid(row=1, column=1, padx=5, sticky=(tk.W, tk.E, tk.N, tk.S))
        trans_scrollbar = ttk.Scrollbar(preview_frame, orient=tk.VERTICAL, command=self.trans_preview.yview)
        trans_scrollbar.grid(row=1, column=1, sticky=(tk.N, tk.S, tk.E))
        self.trans_preview.configure(yscrollcommand=trans_scrollbar.set)
        self.trans_preview.configure(state='disabled')
        
        # 设置主框架的行权重，使预览区域可以扩展
        main_frame.rowconfigure(7, weight=1)
        
        # 设置窗口最小尺寸
        root.minsize(800, 600)

    def update_srt_preview(self, filename):
        """更新SRT文件预览"""
        self.srt_preview.configure(state='normal')
        self.srt_preview.delete(1.0, tk.END)
        try:
            with open(filename, 'r', encoding='utf-8') as f:
                content = f.read()
                # 只显示前10个字幕块或2000个字符
                blocks = content.split('\n\n')[:10]
                preview_content = '\n\n'.join(blocks)
                if len(preview_content) > 2000:
                    preview_content = preview_content[:2000] + "...\n(显示部分内容)"
                self.srt_preview.insert(tk.END, preview_content)
        except Exception as e:
            self.srt_preview.insert(tk.END, f"无法读取文件: {str(e)}")
        self.srt_preview.configure(state='disabled')

    def update_trans_preview(self, filename):
        """更新翻译文件预览"""
        self.trans_preview.configure(state='normal')
        self.trans_preview.delete(1.0, tk.END)
        try:
            with open(filename, 'r', encoding='utf-8') as f:
                lines = f.readlines()
                # 只显示前20行或2000个字符
                preview_lines = lines[:20]
                preview_content = ''.join(preview_lines)
                if len(preview_content) > 2000:
                    preview_content = preview_content[:2000] + "...\n(显示部分内容)"
                self.trans_preview.insert(tk.END, preview_content)
        except Exception as e:
            self.trans_preview.insert(tk.END, f"无法读取文件: {str(e)}")
        self.trans_preview.configure(state='disabled')

    def select_srt(self):
        filename = filedialog.askopenfilename(filetypes=[("SRT文件", "*.srt")])
        if filename:
            self.srt_path.set(filename)
            output_dir = os.path.dirname(filename)
            base_name = os.path.splitext(os.path.basename(filename))[0]
            self.output_path.set(os.path.join(output_dir, f"{base_name}_translated.srt"))
            # 更新SRT行数统计
            srt_count = count_srt_lines(filename)
            self.srt_count_var.set(f"SRT字幕数: {srt_count}")
            # 更新预览
            self.update_srt_preview(filename)

    def select_translation(self):
        filename = filedialog.askopenfilename(filetypes=[("文本文件", "*.txt")])
        if filename:
            self.trans_path.set(filename)
            # 更新TXT行数统计
            txt_count = count_txt_lines(filename)
            self.txt_count_var.set(f"TXT行数: {txt_count}")
            # 更新预览
            self.update_trans_preview(filename)

    def select_output(self):
        filename = filedialog.asksaveasfilename(
            defaultextension=".srt",
            filetypes=[("SRT文件", "*.srt")]
        )
        if filename:
            self.output_path.set(filename)

    def convert(self):
        try:
            if not self.srt_path.get() or not self.trans_path.get() or not self.output_path.get():
                raise ValueError("请选择所有必需的文件！")

            # 检查行数是否匹配
            srt_count = count_srt_lines(self.srt_path.get())
            txt_count = count_txt_lines(self.trans_path.get())
            
            if srt_count != txt_count:
                raise ValueError(f"字幕数量与翻译数量不匹配！\nSRT字幕数: {srt_count}\nTXT行数: {txt_count}")

            self.status_var.set("正在转换...")
            self.root.update()

            subtitles = read_srt(self.srt_path.get())
            translations = read_translation(self.trans_path.get())
            
            # 获取输出文件路径
            srt_output = self.output_path.get()
            txt_output = None
            
            if self.export_txt.get():
                # 生成TXT文件路径（与SRT文件同名，但扩展名为.txt）
                txt_output = os.path.splitext(srt_output)[0] + "_对照.txt"
            
            # 创建输出文件
            create_new_srt(subtitles, translations, srt_output, txt_output)

            success_msg = "转换完成！"
            if txt_output:
                success_msg += f"\nSRT文件：{srt_output}\nTXT对照文件：{txt_output}"
            
            self.status_var.set("转换完成！")
            messagebox.showinfo("成功", success_msg)
        
        except Exception as e:
            self.status_var.set(f"错误: {str(e)}")
            messagebox.showerror("错误", str(e))

def main():
    root = tk.Tk()
    # 尝试设置窗口样式（在Windows上可用）
    try:
        from ctypes import windll
        windll.shcore.SetProcessDpiAwareness(1)
    except:
        pass
    
    app = SrtTranslatorGUI(root)
    root.mainloop()

if __name__ == "__main__":
    main() 